{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import necessary modules\n",
    "# uncomment to get plots displayed in notebook\n",
    "%matplotlib inline\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from classy import Class\n",
    "from scipy.optimize import fsolve\n",
    "from scipy.interpolate import interp1d\n",
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# esthetic definitions for the plots\n",
    "font = {'size'   : 16, 'family':'STIXGeneral'}\n",
    "axislabelfontsize='large'\n",
    "matplotlib.rc('font', **font)\n",
    "matplotlib.mathtext.rcParams['legend.fontsize']='medium'\n",
    "plt.rcParams[\"figure.figsize\"] = [8.0,6.0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#############################################\n",
    "#\n",
    "# value of k that we want to follow in [1/Mpc]\n",
    "#\n",
    "k = 0.5  # 1/Mpc\n",
    "#\n",
    "# Cosmological parameters and other CLASS parameters\n",
    "#\n",
    "common_settings = {# we need to set the output field to something although\n",
    "                   # the really releveant outpout here will be set with 'k_output_values'\n",
    "                   'output':'mPk',\n",
    "                   # value of k we want to polot in [1/Mpc]\n",
    "                   'k_output_values':k,\n",
    "                   # LambdaCDM parameters\n",
    "                   'h':0.67810,\n",
    "                   'omega_b':0.02238280,\n",
    "                   'omega_cdm':0.1201075,\n",
    "                   'A_s':2.100549e-09 ,\n",
    "                   'n_s':0.9660499,\n",
    "                   'tau_reio':0.05430842,\n",
    "                   # Take fixed value for primordial Helium (instead of automatic BBN adjustment)\n",
    "                   'YHe':0.2454,\n",
    "                   # other options and settings\n",
    "                   'compute damping scale':'yes', # needed to output the time of damping scale crossing\n",
    "                   'gauge':'newtonian'}  \n",
    "##############\n",
    "#    \n",
    "# call CLASS\n",
    "#\n",
    "M = Class()\n",
    "M.set(common_settings)\n",
    "M.compute()\n",
    "#\n",
    "# load perturbations\n",
    "#\n",
    "all_k = M.get_perturbations()  # this potentially constains scalars/tensors and all k values\n",
    "print (all_k['scalar'][0].keys())\n",
    "#    \n",
    "one_k = all_k['scalar'][0]     # this contains only the scalar perturbations for the requested k values\n",
    "#    \n",
    "tau = one_k['tau [Mpc]']\n",
    "Theta0 = 0.25*one_k['delta_g']\n",
    "phi = one_k['phi']\n",
    "psi = one_k['psi']\n",
    "theta_b = one_k['theta_b']\n",
    "a = one_k['a']\n",
    "# compute related quantitites    \n",
    "R = 3./4.*M.Omega_b()/M.Omega_g()*a    # R = 3/4 * (rho_b/rho_gamma)\n",
    "zero_point = -(1.+R)*psi               # zero point of oscillations: -(1.+R)*psi\n",
    "#\n",
    "# get Theta0 oscillation amplitude (for vertical scale of plot)\n",
    "#\n",
    "Theta0_amp = max(Theta0.max(),-Theta0.min())\n",
    "#\n",
    "# get the time of decoupling\n",
    "#\n",
    "quantities = M.get_current_derived_parameters(['tau_rec'])\n",
    "# print times.viewkeys()\n",
    "tau_rec = quantities['tau_rec']\n",
    "#\n",
    "# use table of background quantitites to find the time of\n",
    "# Hubble crossing (k / (aH)= 2 pi), sound horizon crossing (k * rs = 2pi)\n",
    "#\n",
    "background = M.get_background() # load background table\n",
    "#print background.viewkeys()\n",
    "#\n",
    "background_tau = background['conf. time [Mpc]'] # read confromal times in background table\n",
    "background_z = background['z'] # read redshift\n",
    "background_k_over_aH = k/background['H [1/Mpc]']*(1.+background['z']) # read k/aH = k(1+z)/H\n",
    "background_k_rs = k * background['comov.snd.hrz.'] # read k * rs\n",
    "background_rho_m_over_r =\\\n",
    "    (background['(.)rho_b']+background['(.)rho_cdm'])\\\n",
    "    /(background['(.)rho_g']+background['(.)rho_ur']) # read rho_r / rho_m (to find time of equality)\n",
    "#\n",
    "# define interpolation functions; we want the value of tau when the argument is equal to 2pi (or 1 for equality)\n",
    "#\n",
    "tau_at_k_over_aH = interp1d(background_k_over_aH,background_tau)\n",
    "tau_at_k_rs = interp1d(background_k_rs,background_tau)\n",
    "tau_at_rho_m_over_r = interp1d(background_rho_m_over_r,background_tau)\n",
    "#\n",
    "# finally get these times\n",
    "#\n",
    "tau_Hubble = tau_at_k_over_aH(2.*math.pi)\n",
    "tau_s = tau_at_k_rs(2.*math.pi)\n",
    "tau_eq = tau_at_rho_m_over_r(1.)\n",
    "#\n",
    "#################\n",
    "#\n",
    "# start plotting\n",
    "#\n",
    "#################\n",
    "#    \n",
    "plt.xlim([tau[0],tau_rec*1.3])\n",
    "plt.ylim([-1.3*Theta0_amp,1.3*Theta0_amp])\n",
    "plt.xlabel(r'$\\tau \\,\\,\\, \\mathrm{[Mpc]}$')\n",
    "plt.title(r'$\\mathrm{Transfer} (\\tau,k) \\,\\,\\, \\mathrm{for} \\,\\,\\, k=%g \\,\\,\\, [1/\\mathrm{Mpc}]$'%k)\n",
    "plt.grid()\n",
    "#\n",
    "plt.axvline(x=tau_Hubble,color='r')\n",
    "plt.axvline(x=tau_s,color='y')\n",
    "plt.axvline(x=tau_eq,color='k')\n",
    "plt.axvline(x=tau_rec,color='k')\n",
    "#\n",
    "plt.annotate(r'Hubble cross.',\n",
    "                xy=(tau_Hubble,1.08*Theta0_amp),\n",
    "                xytext=(0.15*tau_Hubble,1.18*Theta0_amp),\n",
    "                arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "plt.annotate(r'sound hor. cross.',\n",
    "                 xy=(tau_s,-1.0*Theta0_amp),\n",
    "                 xytext=(1.5*tau_s,-1.2*Theta0_amp),\n",
    "                 arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "plt.annotate(r'eq.',\n",
    "                 xy=(tau_eq,1.08*Theta0_amp),\n",
    "                 xytext=(0.45*tau_eq,1.18*Theta0_amp),\n",
    "                 arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "plt.annotate(r'rec.',\n",
    "                 xy=(tau_rec,1.08*Theta0_amp),\n",
    "                 xytext=(0.45*tau_rec,1.18*Theta0_amp),\n",
    "                 arrowprops=dict(facecolor='black', shrink=0.05, width=1, headlength=5, headwidth=5))\n",
    "#\n",
    "# Possibility to add functions one by one, saving between each (for slides)\n",
    "#\n",
    "plt.semilogx(tau,psi,'y-',label=r'$\\psi$')\n",
    "#plt.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "#plt.savefig('one_k_1.pdf',bbox_inches='tight')\n",
    "#\n",
    "plt.semilogx(tau,phi,'r-',label=r'$\\phi$')\n",
    "#plt.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "#plt.savefig('one_k_2.pdf',bbox_inches='tight')\n",
    "#\n",
    "plt.semilogx(tau,zero_point,'k:',label=r'$-(1+R)\\psi$')\n",
    "#plt.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "#plt.savefig('one_k_3.pdf',bbox_inches='tight')\n",
    "#\n",
    "plt.semilogx(tau,Theta0,'b-',linewidth=2,label=r'$\\Theta_0$')\n",
    "#plt.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "#plt.savefig('one_k_4.pdf',bbox_inches='tight')\n",
    "#\n",
    "plt.semilogx(tau,Theta0+psi,'c-',linewidth=2,label=r'$\\Theta_0+\\psi$')\n",
    "#plt.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "#plt.savefig('one_k_5.pdf',bbox_inches='tight')\n",
    "#\n",
    "plt.semilogx(tau,theta_b,'g-',label=r'$\\theta_b$')\n",
    "plt.legend(loc='right',bbox_to_anchor=(1.4, 0.5))\n",
    "plt.savefig('one_k.pdf',bbox_inches='tight')\n",
    "#"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
